{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b50c4af8-fec3-4396-860a-1322089d76cb",
   "metadata": {},
   "source": [
    "# OpenAI Agent with Query Engine Tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db402a8b-90d6-4e1d-8df6-347c54624f26",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Build Query Engine Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "02160804-64a2-4ef3-8a0d-8c16b06fd205",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/suo/miniconda3/envs/llama/lib/python3.9/site-packages/deeplake/util/check_latest_version.py:32: UserWarning: A newer version of deeplake (3.6.7) is available. It's recommended that you update to the latest version using `pip install -U deeplake`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from llama_index import (\n",
    "    SimpleDirectoryReader,\n",
    "    VectorStoreIndex,\n",
    "    StorageContext,\n",
    "    load_index_from_storage,\n",
    ")\n",
    "\n",
    "from llama_index.tools import QueryEngineTool, ToolMetadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "91618236-54d3-4783-86b7-7b7554efeed1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    storage_context = StorageContext.from_defaults(persist_dir=\"./storage/lyft\")\n",
    "    lyft_index = load_index_from_storage(storage_context)\n",
    "\n",
    "    storage_context = StorageContext.from_defaults(persist_dir=\"./storage/uber\")\n",
    "    uber_index = load_index_from_storage(storage_context)\n",
    "\n",
    "    index_loaded = True\n",
    "except:\n",
    "    index_loaded = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d3d0bb8c-16c8-4946-a9d8-59528cf3952a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if not index_loaded:\n",
    "    # load data\n",
    "    lyft_docs = SimpleDirectoryReader(\n",
    "        input_files=[\"../data/10k/lyft_2021.pdf\"]\n",
    "    ).load_data()\n",
    "    uber_docs = SimpleDirectoryReader(\n",
    "        input_files=[\"../data/10k/uber_2021.pdf\"]\n",
    "    ).load_data()\n",
    "\n",
    "    # build index\n",
    "    lyft_index = VectorStoreIndex.from_documents(lyft_docs)\n",
    "    uber_index = VectorStoreIndex.from_documents(uber_docs)\n",
    "\n",
    "    # persist index\n",
    "    lyft_index.storage_context.persist(persist_dir=\"./storage/lyft\")\n",
    "    uber_index.storage_context.persist(persist_dir=\"./storage/uber\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "31892898-a2dc-43c8-812a-3442feb2108d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "lyft_engine = lyft_index.as_query_engine(similarity_top_k=3)\n",
    "uber_engine = uber_index.as_query_engine(similarity_top_k=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f9f3158a-7647-4442-8de1-4db80723b4d2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "query_engine_tools = [\n",
    "    QueryEngineTool(\n",
    "        query_engine=lyft_engine,\n",
    "        metadata=ToolMetadata(\n",
    "            name=\"lyft_10k\",\n",
    "            description=\"Provides information about Lyft financials for year 2021. \"\n",
    "            \"Use a detailed plain text question as input to the tool.\",\n",
    "        ),\n",
    "    ),\n",
    "    QueryEngineTool(\n",
    "        query_engine=uber_engine,\n",
    "        metadata=ToolMetadata(\n",
    "            name=\"uber_10k\",\n",
    "            description=\"Provides information about Uber financials for year 2021. \"\n",
    "            \"Use a detailed plain text question as input to the tool.\",\n",
    "        ),\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "275c01b1-8dce-4216-9203-1e961b7fc313",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Setup OpenAI Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "32f71a46-bdf6-4365-b1f1-e23a0d913a3d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.agent import OpenAIAgent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ded93297-fee8-4329-bf37-cf77e87621ae",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "agent = OpenAIAgent.from_tools(query_engine_tools, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22716961-11c3-4ac4-82a8-419f787bc36a",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Let's Try It Out!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42a1bce0-b398-4937-9008-6cee04368ac4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "===== Entering Chat REPL =====\n",
      "Type \"exit\" to exit.\n",
      "\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "Human:  Compare and contrast Lyft and Uber's revenue growth in 2021, then give an analysis\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Calling Function ===\n",
      "Calling function: lyft_10k with args: {\n",
      "  \"input\": \"What was Lyft's revenue growth in 2021?\"\n",
      "}\n",
      "Got output: \n",
      "Lyft's revenue growth in 2021 was 36%.\n",
      "========================\n",
      "=== Calling Function ===\n",
      "Calling function: uber_10k with args: \n",
      "{\n",
      "  \"input\": \"What was Uber's revenue growth in 2021?\"\n",
      "}\n",
      "Got output: \n",
      "Uber's revenue growth in 2021 was 57%.\n",
      "========================\n",
      "Assistant: In 2021, both Lyft and Uber experienced significant revenue growth. Lyft's revenue grew by 36%, while Uber's revenue grew by 57%. \n",
      "\n",
      "The higher revenue growth rate of Uber indicates that it had a stronger performance in terms of generating revenue compared to Lyft. This could be attributed to several factors, including Uber's larger market presence and global reach. Uber operates in more countries and cities compared to Lyft, which allows it to capture a larger customer base and generate more revenue.\n",
      "\n",
      "However, it's important to note that revenue growth alone does not provide a complete picture of a company's financial performance. Other factors such as profitability, market share, and operational efficiency also play a crucial role in assessing the overall success of a company.\n",
      "\n",
      "In terms of revenue growth, Uber outperformed Lyft in 2021. However, a comprehensive analysis would require considering other financial metrics and factors to get a complete understanding of the two companies' performance in the market.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "agent.chat_repl()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0873463d-4790-4c17-bfe9-2ece610fe4b3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
